﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

namespace V32Assembler
{
    public static class Assembler
    {
        private static List<Output> Outputs { get; set; }

        public static List<Error> Errors { get; private set; }
        public static List<Warning> Warnings { get; private set; }
        public static int CurrentLineNumber { get; private set; }

        private static List<Macro> Macros { get; set; }
        public static void AddMacro(Macro macro)
        {
            foreach (Macro existingMacro in Macros)
                if (existingMacro.Name == macro.Name
                    && existingMacro.ArgNames.Length == macro.ArgNames.Length)
                    throw new Error("Invalid overload of macro '" + macro.Name + "'. Macros must be overloaded by argument count.");

            Macros.Add(macro);
        }

        private static List<Tuple<string, int>> filterCode(string[] program)
        {
            List<Tuple<string, int>> programList = new List<Tuple<string, int>>();
            for (int i = 0; i < program.Length; i++)
            {
                if (program[i].Contains("#"))
                    program[i] = program[i].Remove(program[i].IndexOf("#"));
                program[i] = program[i].Trim();
                if (program[i].Length != 0)
                    programList.Add(new Tuple<string, int>(program[i], i));
            }
            return programList;
        }

        private static Dictionary<string, int> getLabelLines(List<Tuple<string, int>> program)
        {
            Dictionary<string, int> labelLines = new Dictionary<string, int>();
            for (int i = 0; i < program.Count; i++)
            {
                try
                {
                    if (program[i].Item1.EndsWith(":"))
                    {
                        string labelName = program[i].Item1.Remove(program[i].Item1.Length - 1).TrimEnd();
                        if (labelLines.ContainsKey(labelName))
                            throw new Error("Label '" + labelName + "' is multiply defined.");
                        labelLines.Add(labelName, i);
                        program.RemoveAt(i--);
                    }
                    //if (program[i].Item1.Contains(":"))
                    //{
                    //    string labelName = program[i].Item1.Remove(program[i].Item1.IndexOf(':')).TrimEnd();
                    //    if (labelLines.ContainsKey(labelName))
                    //        throw new Error("Label '" + labelName + "' is multiply defined.");
                    //    labelLines.Add(labelName, i);

                    //    string restOfLine = program[i].Item1.Substring(program[i].Item1.IndexOf(':') + 1).Trim();
                    //    if (restOfLine == "")
                    //        program.RemoveAt(i--);
                    //    else program[i] = new Tuple<string,int>(restOfLine, program[i].Item2);
                    //}
                }
                catch (Error) { }
            }
            return labelLines;
        }

        private static List<Tuple<string, int>> includeFiles(List<Tuple<string, int>> program)
        {
            for (int i = 0; i < program.Count; i++)
            {
                Tuple<string, int> programLine = program[i];
                string line = programLine.Item1;
                int lineNum = CurrentLineNumber = programLine.Item2;

                if (line.ToLower().StartsWith(".include "))
                {
                    string fileName = line.Substring(".include ".Length).Trim().ParseString();
                    // need to make relative to the file, not the executable
                    if (!System.IO.File.Exists(fileName)) throw new Error("File '" + fileName + "' could not be found.");
                    try
                    {
                        System.IO.FileInfo info = new System.IO.FileInfo(fileName);
                        string[] includedLines = System.IO.File.ReadAllLines(fileName);

                        program.RemoveAt(i);
                        foreach (string includedLine in includedLines)
                            program.Insert(i++, new Tuple<string, int>(includedLine, lineNum));
                        i--;
                    }
                    catch (Exception e)
                    {
                        throw new Error("File '" + fileName + "' could not be opened due to error: \"" + e.Message + "\".");
                    }
                }
            }
            return program;
        }

        private static List<Tuple<string, int>> expandMacros(List<Tuple<string, int>> program)
        {
            List<Tuple<string, int>> programLines = new List<Tuple<string,int>>();
            foreach (Tuple<string, int> programLine in program) programLines.Add(programLine);

            // find all macro definitions.
            for (int i = 0; i < program.Count; )
            {
                Tuple<string, int> programLine = program[i];
                string line = programLine.Item1;
                int lineNum = CurrentLineNumber = programLine.Item2;

                if (line.ToLower().StartsWith(".define "))
                {
                    string[] args = line.Substring(".define ".Length).Split(new char[] { ' ' }, StringSplitOptions.RemoveEmptyEntries);
                    if (args.Length < 1) throw new Error("The name of a macro must come after a '.define'.");
                    string name = args[0];
                    OpCode nameAsOpcode;
                    if (Enum.TryParse<OpCode>(name, true, out nameAsOpcode))
                        throw new Error("The name of a macro must not be the name of an opcode.");
                    List<string> argList = new List<string>();
                    for (int argIndex = 1; argIndex < args.Length; argIndex++)
                    {
                        string arg = args[argIndex];
                        if (Register.FromString(arg) != Register.None)
                            throw new Error("The name of a macro argument may not be the name of a register.");
                        if (argList.Contains(arg)) throw new Error("A macro may not have multiple arguments with the same name.");
                        argList.Add(arg);
                    }
                    args = argList.ToArray();

                    program.RemoveAt(i);
                    if (i >= program.Count) throw new Error("A '.define' must be followed by a macro definition.");

                    List<string> macroCommands = new List<string>();
                    for (; i < program.Count && program[i].Item1.ToLower() != ".end"; program.RemoveAt(i))
                        macroCommands.Add(program[i].Item1);
                    if (i >= program.Count) throw new Error("A macro definition must end with a '.end'.");
                    else program.RemoveAt(i);

                    AddMacro(new Macro(name, args, macroCommands.ToArray()));
                }
                else i++;
            }

            List<Tuple<string, int>> temp = program;
            program = programLines;
            programLines = temp;

            // keep expanding the macros until there is no difference from before.
            while (true)
            {
                temp = programLines.ShallowCopy();
                programLines = expandMacros(program, programLines);
                // if temp and programLines are the same, then exit.
                if (temp.Count == programLines.Count)
                {
                    bool same = true;
                    for (int i = 0; i < temp.Count; i++)
                        if (temp[i] != programLines[i])
                            same = false;
                    if (same)
                        return programLines;
                }
            }
        }
        private static List<Tuple<string, int>> expandMacros(List<Tuple<string, int>> program, List<Tuple<string, int>> programLines)
        {
            // actually expand macros here
            for (int i = 0; i < programLines.Count; i++)
            {
                Tuple<string, int> programLine = programLines[i];
                string line = programLine.Item1;
                int lineNum = CurrentLineNumber = programLine.Item2;

                string command = line.Contains(' ') ? line.Remove(line.IndexOf(' ')) : line;
                List<Macro> macros = new List<Macro>();
                foreach (Macro macro in Macros)
                    if (macro.Name == command)
                        macros.Add(macro);
                string[] args = line.Contains(' ') ? line.Remove(0, line.IndexOf(' ')).Split(new char[] { ' ' }, StringSplitOptions.RemoveEmptyEntries) : new string[0];
                foreach (Macro macro in macros)
                {
                    if (args.Length == macro.ArgNames.Length)
                    {
                        programLines.RemoveAt(i);
                        string[] expansion = macro.Expand(args);
                        foreach (string expansionLine in expansion)
                            programLines.Insert(i++, new Tuple<string, int>(expansionLine, lineNum));
                        i--;
                    }
                }
            }
            return programLines;
        }

        private static Tuple<Dictionary<PartialCommand, uint>, List<int>> getCommandAddresses(List<Tuple<string, int>> program, Dictionary<string, int> labelLines)
        {
            Dictionary<PartialCommand, uint> addresses = new Dictionary<PartialCommand, uint>();
            List<int> errorLines = new List<int>();

            uint ip = 0;
            int programIndex = 0;
            foreach (Tuple<string, int> line in program)
            {
                try
                {
                    CurrentLineNumber = line.Item2;
                    PartialCommand command = Command.FromLine(line.Item1, labelLines);
                    if (command == null)
                    {
                        errorLines.Add(programIndex++);
                        continue;
                    }
                    addresses.Add(command, ip);
                    ip += (uint)command.Command.ToBytes().Length;
                    programIndex++;
                }
                catch (Error) { errorLines.Add(programIndex++); }// e) { Assembler.Errors.Add(e); }
                catch (Exception e)
                {
                    if (e.Message == ".datasect")
                        continue;
                    else throw;
                }
            }

            return new Tuple<Dictionary<PartialCommand,uint>,List<int>>(addresses, errorLines);
        }
        private static Dictionary<string, uint> getLabelAddresses(List<Tuple<string, PartialCommand>> program,
            Dictionary<string, int> labelLines,
            Tuple<Dictionary<PartialCommand, uint>, List<int>> commandAddresses)
        {
            try
            {
                Dictionary<string, uint> labelAddresses = new Dictionary<string, uint>();

                foreach (Tuple<string, PartialCommand> lineAndCommand in program)
                {
                    string line = lineAndCommand.Item1;
                    PartialCommand command = lineAndCommand.Item2;

                    if (command.Arg1Label == "" && command.Arg2Label == "") continue; // this command doesn't use any labels

                    if (command.Arg1Label != "")
                    {
                        int labelLineNum = labelLines[command.Arg1Label];
                        foreach (int errorLine in commandAddresses.Item2)
                            if (labelLineNum >= errorLine)
                                labelLineNum--;

                        PartialCommand labelCommand = labelLineNum == program.Count ? PartialCommand.Null : program[labelLineNum].Item2;

                        uint commandAddress = commandAddresses.Item1[labelCommand];
                        if (!labelAddresses.ContainsKey(command.Arg1Label))
                            labelAddresses.Add(command.Arg1Label, commandAddress);
                    }
                    if (command.Arg2Label != "")
                    {
                        int labelLineNum = labelLines[command.Arg2Label];
                        foreach (int errorLine in commandAddresses.Item2)
                            if (labelLineNum >= errorLine)
                                labelLineNum--;

                        PartialCommand labelCommand = labelLineNum == program.Count ? PartialCommand.Null : program[labelLineNum].Item2;

                        uint commandAddress = commandAddresses.Item1[labelCommand];
                        if (!labelAddresses.ContainsKey(command.Arg2Label))
                            labelAddresses.Add(command.Arg2Label, commandAddress);
                    }
                }

                return labelAddresses;
            }
            catch (Exception)
            {
                return null;
            }
        }

        /* object file:
         * ------------
         * V32OBJ
         * [NumLabels]([LabelLength][Label][LabelAddress] [NumUsesOfLabel]([UseAddress]*))*
         * [Program]
         */
        private static byte[] getObjectFileBytes(List<PartialCommand> commands, Dictionary<string, uint> labelAddresses, Dictionary<PartialCommand, uint> commandAddresses)
        {
            List<byte> bytes = new List<byte>();
            System.Text.Encoding strEncoding = new System.Text.ASCIIEncoding();

            string magicHeader = "V32OBJ";
            bytes.AddRange(strEncoding.GetBytes(magicHeader)); // write header

            bytes.AddRange(((uint)labelAddresses.Count).ToBytes()); // write num labels
            // we want to write each label in order of definition, not in order of use. so sort by the labels' addresses.
            foreach (KeyValuePair<string, uint> labelAddress in (from kvp in labelAddresses orderby kvp.Value ascending select kvp))
            {
                string label = labelAddress.Key;
                uint address = labelAddress.Value;

                bytes.AddRange(((uint)strEncoding.GetByteCount(label)).ToBytes()); // write label length
                bytes.AddRange(strEncoding.GetBytes(label)); // write label
                bytes.AddRange(((uint)address).ToBytes()); // write label address

                List<uint> labelUses = new List<uint>();
                foreach (PartialCommand command in commands)
                {
                    //if (command.Command is DataCommand || command.Command is DataBlockCommand) continue;

                    uint commandAddress = commandAddresses[command];
                    uint arg1Address = commandAddress + command.Command.GetArg1Offset(),
                         arg2Address = commandAddress + command.Command.GetArg2Offset();

                    if (command.Arg1Label == label)
                        labelUses.Add(arg1Address); // write address of the use of the label
                    if (command.Arg2Label == label)
                        labelUses.Add(arg2Address); // write address of the use of the label
                }
                bytes.AddRange(((uint)labelUses.Count).ToBytes()); // write the number of uses of this label
                foreach (uint labelUse in labelUses)
                    bytes.AddRange(labelUse.ToBytes()); // write the addresses of the uses of this label
            }

            // write compiled commands
            foreach (PartialCommand commandAndLabels in commands)
            {
                Command command = commandAndLabels.Command;
                bytes.AddRange(command.ToBytes());
            }

            return bytes.ToArray();
        }

        public static Output[] Assemble(string[] input, bool preprocess, bool addresses)
        {
            Outputs = new List<Output>();

            Errors = new List<Error>();
            Warnings = new List<Warning>();
            Macros = new List<Macro>();

            List<Tuple<string, int>> program = filterCode(input); // list of tuple of line of code (minus comments, but including errors) and actual line numbers
            program = includeFiles(program);
            program = expandMacros(program);

            if (preprocess)
            {
                string[] lines = (from line in program select line.Item1).ToArray();
                for (int i = 0; i < lines.Length; i++)
                {
                    if (!lines[i].TrimEnd().EndsWith(":"))
                        lines[i] = "    " + lines[i];
                }
                Outputs.Add(new PreprocessedFileOutput(lines));
            }

            Dictionary<string, int> labelLines = getLabelLines(program); // dictionary from label to line of code (including errors and macro expansions)

            Tuple<Dictionary<PartialCommand, uint>, List<int>> commandAddresses = getCommandAddresses(program, labelLines);
            

            List<Tuple<string, PartialCommand>> programCommands = new List<Tuple<string,PartialCommand>>();
            int lineNumber = 0;
            Tuple<PartialCommand, uint> highestCommandAddress = new Tuple<PartialCommand, uint>(null, 0);
            foreach (PartialCommand command in commandAddresses.Item1.Keys) // assuming that Keys is in the order that KVP's were added
            {
                if (commandAddresses.Item1[command] > highestCommandAddress.Item2)
                    highestCommandAddress = new Tuple<PartialCommand,uint>(command, commandAddresses.Item1[command]);
                while (commandAddresses.Item2.Contains(lineNumber)) lineNumber++;
                string line = program[lineNumber].Item1;
                programCommands.Add(new Tuple<string,PartialCommand>(line, command));
                lineNumber++;
            }
            commandAddresses.Item1.Add(PartialCommand.Null, highestCommandAddress.Item2 + (uint)highestCommandAddress.Item1.Command.ToBytes().Length);

            Dictionary<string, uint> labelAddresses = getLabelAddresses(programCommands, labelLines, commandAddresses);
            if (labelAddresses == null)
                return null;

            if (addresses)
            {
                if (Errors.Count > 0) return null;

                List<string> lines = new List<string>();
                foreach (Tuple<string, PartialCommand> programCommand in programCommands)
                {
                    uint commandAddress = commandAddresses.Item1[programCommand.Item2];
                    List<string> labels = new List<string>();
                    foreach (KeyValuePair<string, uint> labelAddress in labelAddresses)
                        if (labelAddress.Value == commandAddress)
                            labels.Add(labelAddress.Key + ":");

                    string line = string.Format("{0:0000000000}:", commandAddress);
                    if (labels.Count == 0) lines.Add(line + "    " + programCommand.Item1);
                    else
                    {
                        lines.Add(line + labels[0]);
                        for (int i = 1; i < labels.Count; i++)
                            lines.Add("----------:" + labels[i]);
                        lines.Add("----------:" + "    " + programCommand.Item1);
                    }
                }

                Outputs.Add(new CommandAddressesFileOutput(lines.ToArray()));
            }

            List<PartialCommand> commands = new List<PartialCommand>();
            foreach (Tuple<string, PartialCommand> command in programCommands)
                commands.Add(command.Item2);

            Outputs.Add(new ObjectFileOutput(getObjectFileBytes(commands, labelAddresses, commandAddresses.Item1)));

            return Outputs.ToArray();
        }
    }
}
